from ast import Not
import stat
from nnunetv2.training.nnUNetTrainer.variants.optimizer.nnUNetTrainerAdam import nnUNetTrainerAdam3en4
from nnunetv2.training.nnUNetTrainer.nnUNetTrainerDist import nnUNetTrainerDist
from nnunetv2.training.loss.deep_supervision import DeepSupervisionWrapper
import torch
import warnings
import numpy as np
from torch import autocast
from nnunetv2.utilities.helpers import empty_cache, dummy_context
from nnunetv2.training.loss.dice import MemoryEfficientSoftDiceLoss, get_tp_fp_fn_tn
from monai.networks import one_hot
# from nnunetv2.training.loss.ce import RobustCrossEntropyLoss

class NeuronSegBaseTrainer(nnUNetTrainerAdam3en4):
    pass

class NeuronSegBaseTrainerNoDeepSupervision(nnUNetTrainerAdam3en4):
    def __init__(
        self,
        plans: dict,
        configuration: str,
        fold: int,
        dataset_json: dict,
        device: torch.device = torch.device("cuda"),
    ):
        super().__init__(plans, configuration, fold, dataset_json, device)
        self.enable_deep_supervision = False
        
    def set_deep_supervision_enabled(self, enabled: bool):
        pass

class nnUNetTrainerStage4(NeuronSegBaseTrainer):
    @staticmethod
    def build_network_architecture(
        architecture_class_name, 
        arch_init_kwargs, 
        arch_init_kwargs_req_import, 
        num_input_channels, 
        num_output_channels, 
        enable_deep_supervision = True):

        arch_init_kwargs['n_stages'] = 4
        arch_init_kwargs['features_per_stage'] = arch_init_kwargs['features_per_stage'][:4]
        arch_init_kwargs['kernel_sizes'] = arch_init_kwargs['kernel_sizes'][:4]
        arch_init_kwargs['strides'] = arch_init_kwargs['strides'][:4]
        arch_init_kwargs['n_conv_per_stage'] = arch_init_kwargs['n_conv_per_stage'][:4]
        arch_init_kwargs['n_conv_per_stage_decoder'] = arch_init_kwargs['n_conv_per_stage_decoder'][:4-1]
        return NeuronSegBaseTrainer.build_network_architecture(
            architecture_class_name, 
            arch_init_kwargs, 
            arch_init_kwargs_req_import, 
            num_input_channels, 
            num_output_channels, 
            enable_deep_supervision
        )
    
class nnUNetTrainerDistStage4(nnUNetTrainerStage4, nnUNetTrainerDist):
    @staticmethod
    def build_network_architecture(
        architecture_class_name, 
        arch_init_kwargs, 
        arch_init_kwargs_req_import, 
        num_input_channels, 
        num_output_channels, 
        enable_deep_supervision = True):
        return nnUNetTrainerStage4.build_network_architecture(
            architecture_class_name, 
            arch_init_kwargs, 
            arch_init_kwargs_req_import, 
            num_input_channels, 
            1, 
            enable_deep_supervision
        )

    def _build_loss(self):
        assert len(self.label_manager.foreground_labels) == 1, 'This trainer only supports binary segmentation'
        if self.label_manager.ignore_label is not None:
            raise NotImplementedError('Ignore label is not supported in this trainer')
        loss = torch.nn.BCEWithLogitsLoss()

        if self.enable_deep_supervision:
            deep_supervision_scales = self._get_deep_supervision_scales()

            # we give each output a weight which decreases exponentially (division by 2) as the resolution decreases
            # this gives higher resolution outputs more weight in the loss
            weights = np.array([1 / (2 ** i) for i in range(len(deep_supervision_scales))])
            weights[-1] = 0

            # we don't use the lowest 2 outputs. Normalize weights so that they sum to 1
            weights = weights / weights.sum()
            # now wrap the loss
            loss = DeepSupervisionWrapper(loss, weights)
        return loss


    def train_step(self, batch: dict) -> dict:
        data = batch['data']
        target = batch['target']
        dist = batch['dist']

        # # import napari
        # # viewer = napari.Viewer()
        # # viewer.add_image(data[0].cpu().numpy(), name='data')
        # # viewer.add_image(target[0][0].cpu().numpy(), name='target')
        # # viewer.add_image(skel[0][0].cpu().numpy(), name='skel')
        # # napari.run()

        data = data.to(self.device, non_blocking=True)
        if isinstance(target, list):
            target = [i.to(self.device, non_blocking=True) for i in target]
            dist = [i.unsqueeze(1).to(self.device, non_blocking=True) for i in dist]
        else:
            target = target.to(self.device, non_blocking=True)
            dist = dist.unsqueeze(1).to(self.device, non_blocking=True)

        self.optimizer.zero_grad(set_to_none=True)
        # Autocast can be annoying
        # If the device_type is 'cpu' then it's slow as heck and needs to be disabled.
        # If the device_type is 'mps' then it will complain that mps is not implemented, even if enabled=False is set. Whyyyyyyy. (this is why we don't make use of enabled=False)
        # So autocast will only be active if we have a cuda device.
        with autocast(self.device.type, enabled=True) if self.device.type == 'cuda' else dummy_context():
            output = self.network(data)
            # del data
            l = self.loss(output, dist)

        if self.grad_scaler is not None:
            self.grad_scaler.scale(l).backward()
            self.grad_scaler.unscale_(self.optimizer)
            torch.nn.utils.clip_grad_norm_(self.network.parameters(), 12)
            self.grad_scaler.step(self.optimizer)
            self.grad_scaler.update()
        else:
            l.backward()
            torch.nn.utils.clip_grad_norm_(self.network.parameters(), 12)
            self.optimizer.step()
        return {'loss': l.detach().cpu().numpy()}
    

    def validation_step(self, batch: dict) -> dict:
        data = batch['data']
        target = batch['target']
        dist = batch['dist']

        data = data.to(self.device, non_blocking=True)
        if isinstance(target, list):
            target = [i.to(self.device, non_blocking=True) for i in target]
            dist = [i.unsqueeze(1).to(self.device, non_blocking=True) for i in dist]
        else:
            target = target.to(self.device, non_blocking=True)
            dist = dist.unsqueeze(1).to(self.device, non_blocking=True)

        # Autocast can be annoying
        # If the device_type is 'cpu' then it's slow as heck and needs to be disabled.
        # If the device_type is 'mps' then it will complain that mps is not implemented, even if enabled=False is set. Whyyyyyyy. (this is why we don't make use of enabled=False)
        # So autocast will only be active if we have a cuda device.
        with autocast(self.device.type, enabled=True) if self.device.type == 'cuda' else dummy_context():
            output = self.network(data)
            del data
            l = self.loss(output, dist)

        # we only need the output with the highest output resolution (if DS enabled)
        if self.enable_deep_supervision:
            output = output[0]
            target = target[0]

        # the following is needed for online evaluation. Fake dice (green line)
        axes = [0] + list(range(2, output.ndim))

        # if self.label_manager.has_regions:
        predicted_segmentation_onehot = one_hot((torch.sigmoid(output) > 0.5).long(), num_classes=2, dim=1)
        # else:
        #     # no need for softmax
        #     output_seg = output.argmax(1)[:, None]
        #     predicted_segmentation_onehot = torch.zeros(output.shape, device=output.device, dtype=torch.float32)
        #     predicted_segmentation_onehot.scatter_(1, output_seg, 1)
        #     del output_seg

        if self.label_manager.has_ignore_label:
            if not self.label_manager.has_regions:
                mask = (target != self.label_manager.ignore_label).float()
                # CAREFUL that you don't rely on target after this line!
                target[target == self.label_manager.ignore_label] = 0
            else:
                if target.dtype == torch.bool:
                    mask = ~target[:, -1:]
                else:
                    mask = 1 - target[:, -1:]
                # CAREFUL that you don't rely on target after this line!
                target = target[:, :-1]
        else:
            mask = None

        tp, fp, fn, _ = get_tp_fp_fn_tn(predicted_segmentation_onehot, target, axes=axes, mask=mask)
        tp_hard = tp.detach().cpu().numpy()
        fp_hard = fp.detach().cpu().numpy()
        fn_hard = fn.detach().cpu().numpy()
        if not self.label_manager.has_regions:
            # if we train with regions all segmentation heads predict some kind of foreground. In conventional
            # (softmax training) there needs tobe one output for the background. We are not interested in the
            # background Dice
            # [1:] in order to remove background
            tp_hard = tp_hard[1:]
            fp_hard = fp_hard[1:]
            fn_hard = fn_hard[1:]

        return {'loss': l.detach().cpu().numpy(), 'tp_hard': tp_hard, 'fp_hard': fp_hard, 'fn_hard': fn_hard}